from ultralytics import YOLO
import torch
from multiprocessing import freeze_support

def print_gpu_utilization():
    print("\nGPU使用情况:")
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}:")
        print(f"- 显存使用: {torch.cuda.memory_allocated(i)/1024**3:.2f}GB")
        print(f"- 显存预留: {torch.cuda.memory_reserved(i)/1024**3:.2f}GB")

def train():
    # 检查可用的GPU
    num_gpus = torch.cuda.device_count()
    print(f"Available GPUs: {num_gpus}")
    
    # Load a model
    model = YOLO("yolov8s-cls.pt")  # 使用更大的模型
    
    # 设置CUDA优化
    if torch.cuda.is_available():
        torch.backends.cudnn.benchmark = True
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
        torch.cuda.empty_cache()
    
    print_gpu_utilization()  # 打印初始GPU使用情况
    
    # 训练模型
    results = model.train(
        data="F:/Desktop/wjx_bishe/YOLO_format_cls",  # 数据集路径
        epochs=15,                                     # 训练轮数
        imgsz=256,                                    # 进一步增加图像尺寸
        batch=256,                                    # 显著增加batch size
        device='0',                                   # 使用单GPU
        workers=32,                                   # 增加worker数量以提高CPU利用率
        name='train_run',                             # 实验名称
        verbose=True,                                 # 显示详细信息
        amp=True,                                     # 使用混合精度训练
        cache=True,                                   # 缓存图像到RAM
        optimizer='AdamW',                            # 使用AdamW优化器
        lr0=0.001,                                    # 初始学习率
        weight_decay=0.0005,                         # 权重衰减
        warmup_epochs=3.0,                           # 预热轮数
        label_smoothing=0.1,                         # 标签平滑
        patience=100,                                # 早停耐心值
        augment=True,                                # 启用数据增强
        mixup=0.2,                                   # 增加mixup增强
        mosaic=0.7,                                  # 增加mosaic增强
        degrees=10.0,                                # 添加旋转增强
        translate=0.2,                               # 添加平移增强
        scale=0.2,                                   # 添加缩放增强
        fliplr=0.5,                                  # 水平翻转概率
        copy_paste=0.1,                              # 添加copy-paste增强
        multi_scale=True,                            # 启用多尺度训练
        rect=False,                                  # 关闭矩形训练以允许更多增强
        overlap_mask=True,                           # 允许mask重叠
        mask_ratio=4,                                # mask比例
        dropout=0.2                                  # 添加dropout
    )
    
    print_gpu_utilization()  # 打印训练后GPU使用情况

if __name__ == '__main__':
    freeze_support()
    train()